#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdlib.h>

#define _USE_MATH_DEFINES
#include <math.h>

#include "raylib.h"
#include "pub_lib.h"
#include "FreeImage.h"
#include "image.h"
#include "dpi_algo.h"

BYTE* morph_strel_disk(int r, int* dis);

BYTE* morph_strel(SE_TYPE type, int r, int *xdis, int *ydis)
{
	int size = 9;
	BYTE* se = NULL;

	switch(type) {
		case SE_TYPE_SQUARE:
			*xdis = r / 2;
			*ydis = *xdis;
			size = r * r;
			break;
		case SE_TYPE_VER:
			*xdis = 0;
			*ydis = r / 2;
			size = r;
			break;
		case SE_TYPE_HOR:
			*xdis = r / 2;
			*ydis = 0;
			size = r;
			break;
		case SE_TYPE_DISK:
		default:
			se = morph_strel_disk(r, xdis);
			*ydis = *xdis;
			return se;
	}

	se = malloc(size);
	memset(se, 1, size);
	return se;
}

/* A - B
 * Eroding the image with a square structuring element
 * of size dim * dim filter
 * (default components are all 1 of filter)
 * and saved result to chg
 *
 * ori[h][w] to chg[h][w]
 * */
void erosion_mem_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	int x, y, xdis, ydis;
	BYTE value, fore_color, bc;

	BYTE* se = morph_strel(type, r, &xdis, &ydis);
	bc = 0;
	fore_color = 255;

	//main loop
	int p, q, match, k, index, ii;
	match = 1;
	for(y = 0; y < h; y++){
		for(x = 0; x < w; x++) {
			//judge if center point match
			index = y * w + x;
			value = ori[index];
			if(value) {
				match = 1;
			} else {
				match = 0; //skip this point
				chg[index] = bc;
				continue;
			}

			k = 0;
			for(q = y - ydis; (q <= (y + ydis)) && match; q++) {
				for(p = x - xdis; (p <= (x + xdis)) && match; p++) {
					if(q < 0 || q >= h || p < 0 || p >= w) //match
						continue;
					ii = q * w + p;
					value = ori[ii];
					if(!value) {
						match = 0;
						break;
					}
					k++;
				}//end for q
			}//end for p

			if(match)
				chg[index] = fore_color;
			else
				chg[index] = bc;
		}//end for x
	} //end for y

	free(se);
}

void erosion_image_common(FIBITMAP* ori, FIBITMAP* chg, SE_TYPE type, int r)
{
	int w, h, size;
	w = FreeImage_GetWidth(ori);
	h = FreeImage_GetHeight(ori);
	size = w * h;

	BYTE* ori_mem = malloc(size);
	BYTE* chg_mem = malloc(size);

	init_image_datab(ori, ori_mem);
	erosion_mem_common(ori_mem, chg_mem, w, h, type, r);
	fill_image_data(chg, chg_mem);

	free(ori_mem);
	free(chg_mem);
}


/* A + B
 * Dilation the image with a square structuring element
 * of size dim * dim filter
 * Image and filer: onverlap by at least one element
 * and saved result to chg
 * */
void dilation_mem_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	int x, y, xdis, ydis;
	BYTE value, fore, bc;

	//init
	fore = 255;
	bc = 0;
	BYTE* se = morph_strel(type, r, &xdis, &ydis);

	int p, q, match, k, index, ii;
	match = 1;
	for(y = 0; y < h; y++){
		for(x = 0; x < w; x++) {
			//judge if center point match
			index = y * w + x;
			value = ori[index];
			if(value) {
				match = 1;
				chg[index] = value;
				continue;
			}

			match = 0;
			k = 0;
			for(q = y - ydis; (q <= (y + ydis)) && !match; q++) {
				for(p = x - xdis; (p <= (x + xdis)) && !match; p++) {
					if(q < 0 || q >= h || p < 0 || p >= w) {
						continue;//boundary
					}
					ii = q * w + p;
					value = ori[ii];
					if(value && se[k]) {
						//check se[k] for 4-connectivity only
						match = 1;
						break;
					}
					k++;
				}//end for q
			}//end for p

			if(match)
				chg[index] = fore;
			else
				chg[index] = bc;
		}//end for x
	} //end for y

	free(se);
}

void dilation_image_common(FIBITMAP* ori, FIBITMAP* chg, SE_TYPE type, int r)
{
	int w, h, size;
	w = FreeImage_GetWidth(ori);
	h = FreeImage_GetHeight(ori);
	size = w * h;

	BYTE* ori_mem = malloc(size);
	BYTE* chg_mem = malloc(size);

	init_image_datab(ori, ori_mem);
	dilation_mem_common(ori_mem, chg_mem, w, h, type, r);
	fill_image_data(chg, chg_mem);

	free(ori_mem);
	free(chg_mem);
}

void morph_open_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	BYTE* temp = malloc(w * h);
	erosion_mem_common(ori, temp, w, h, type, r);
	dilation_mem_common(temp, chg, w, h, type, r);
	free(temp);
}

void morph_close_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	BYTE* temp = malloc(w * h);
	dilation_mem_common(ori, temp, w, h, type, r);
	erosion_mem_common(temp, chg, w, h, type, r);
	free(temp);
}

void morph_boundary_extraction(FIBITMAP* ori, FIBITMAP* chg, SE_TYPE type, int r)
{
	//chg = ori - filter
	erosion_image_common(ori, chg, type, r);
	//chg = ori - (ori - filter)
	image_add(ori, chg, chg, -1);
}

//invert
void complement_area(BYTE* buf, int size, int fill)
{
	int i;
	for(i = 0; i < size; i++) {
		buf[i] = fill - buf[i];
	}
}

//get area data from ori
//area: defined by hole rectangle
void image_mem_fill_area(BYTE* ori, int w, int h, morph_label* hole, BYTE* area)
{
	int x, y, index, k;
	k = 0;
	int cw, ch, chg_size;
	cw = hole->ex - hole->sx;
	ch = hole->ey - hole->sy;
	chg_size = cw * ch;

	if(cw == w && ch == h) {
		memcpy(area, ori, chg_size);
		return;
	}

	for(y = hole->sy; y < hole->ey; y++) {
		for(x = hole->sx; x < hole->ex; x++) {
			index = y * w + x;
			area[k] = ori[index];
			k++;
		}
	}
}

//apply changes, write area to chg
//area: defined by hole rectangle
void image_mem_apply_change(BYTE* chg, int w, int h, morph_label* hole, BYTE* area)
{
	int x, y, index, k;
	k = 0;
	int cw, ch, chg_size;
	cw = hole->ex - hole->sx;
	ch = hole->ey - hole->sy;
	chg_size = cw * ch;

	if(cw == w && ch == h) {
		memcpy(chg, area, chg_size);
		return;
	}

	for(y = hole->sy; y < hole->ey; y++) {
		for (x = hole->sx; x < hole->ex; x++) {
			index = y * w + x;
			chg[index] = area[k];
			k++;
		}
	}
}

//result = a intersection b
void morph_sets_intersection(BYTE* a, BYTE* b, BYTE* result, int size)
{
	int i;
	for(i = 0; i < size; i++) {
		if(a[i] && b[i])
			result[i] = a[i];
		else
			result[i] = 0;
	}
}

//result = a union b
void morph_sets_union(BYTE* a, BYTE* b, BYTE* result, int size)
{
	int i;
	for(i = 0; i < size; i++) {
		if(a[i] || b[i])
			result[i] = a[i] + b[i];
		else
			result[i] = 0;
	}
}


/*9.5.2
 * fill a hole of the original image
 * hole position: specified by 'hole' parameter
 *
 * input:
 *  h, w: original image height and width
 *	ori[h][w]: data of original image
 * hole: hole information
 *
 * output:
 *  chg[h][w]
 */
void morph_hole_filling(BYTE* ori, BYTE* chg, int w, int h,
	morph_label* hole, SE_TYPE type, int r, BYTE* label_ary)
{
	int total_size = w * h;
	memcpy(chg, ori, total_size);

	int cw, ch, chg_size;
	int fill = 255;
	cw = hole->ex - hole->sx;
	ch = hole->ey - hole->sy;
	chg_size = cw * ch;
	BYTE* areaA = malloc(chg_size);
	BYTE* xkpre = malloc(chg_size);
	BYTE* xk = malloc(chg_size);
	BYTE* complementA = malloc(chg_size);
	BYTE* tempA = malloc(chg_size);

	//fill areaA
	image_mem_fill_area(ori, w, h, hole, areaA);

	//invert A
	memcpy(complementA, areaA, chg_size);
	complement_mem(complementA, chg_size);

	//init point set to 1
	memset(xkpre, 0, chg_size);
	int index = (hole->k0y - hole->sy) * cw + (hole->k0x - hole->sx);
	xkpre[index] = fill;

	int k = 1;
	do {
		//chg = X(k - 1) + B) and ori^c
		//tempA = X(k - 1) + B
		dilation_mem_common(xkpre, tempA, cw, ch, type, r);
		//xk = tempA intersection complementA = X(k - 1) + B and ori^c
		morph_sets_intersection(tempA, complementA, xk, chg_size);

		//check if equal
		if(memcmp(xkpre, xk, chg_size)) {
			k++;
			memcpy(xkpre, xk, chg_size);
			if(k % 10 == 0)
				printf("background label ongoing %d\n", k);
		} else {
			break;
		}
	}while(k < 1000);
	printf("\nbackground label finished: iteration step k is %d\n", k);

	//label background
	if(label_ary) {
		int i, num;
		memset(label_ary, 0, total_size);
		for(i = 0, num = 0; i < chg_size; i++) {
			if(xk[i] == fill) {
				num++;
				label_ary[i] = MORPH_START_LABEL;
			}
		}
		hole->label = MORPH_START_LABEL;
		hole->pixels = num;
		printf("background pixel num is %d\n\n", num);
	}

	morph_sets_union(xk, areaA, xk, chg_size);
	//apply this change, xk to chg
	image_mem_apply_change(chg, w, h, hole, xk);

	free(areaA);
	free(complementA);
	free(xkpre);
	free(xk);
	free(tempA);
}

//9.5.3 Extraction of Connected Components

/*
 * 1. make label
 * 2. stat pixel number within it
 * 3. calc rectangle of connected components
 *
 * input:
 *	xk[h][w]: include one connected components only
 * */
static void morph_stat_for_con(BYTE* xk, BYTE* label_ary, int w, int h,
	morph_label* label, BYTE fill)
{
	int num = 0, x, y;
	int xmax = 0, ymax = 0, xmin = w, ymin = h;
	int size, index;

	size = w * h;
	for(y = 0; y < h; y++) {
		for (x = 0; x < w; x++) {
			index = y * w + x;
			if (xk[index] == fill) {
				num++;
				label_ary[index] = label->label;
				if (x < xmin)
					xmin = x;
				if (x > xmax)
					xmax = x;
				if (y < ymin)
					ymin = y;
				if (y > ymax)
					ymax = y;
			}
		}
	}

	label->pixels = num;
	label->sx = xmin;
	label->sy = ymin;
	label->ex = xmax;
	label->ey = ymax;

	printf("stat label %d: pixel num is %d\n\n", label->label, num);
}

/*
 * label one connected components
 * morph_label* label: it is in-out parameter
 *
 * input:
 *  h, w: image height and weight
 *  ori[h][w]: image data
 *
 *  label->k0x, label->k0y: coordinate of start point, and
 *    color = ori[k0y][k0x], it is foreground color
 *  label->label: label of this connected components
 *
 * output:
 *  label_ary[h][w]: labeled map
 *
 *  label->pixels: pixel number in this connected components
 *  label->k0x, k0y: init point of hole(with background color)
 *	label->sx, sy, ex, ey: rectangle range of this connected components
 *
 * return value:
 *	k steps
 * */
int morph_label_fore_con(BYTE* ori, BYTE* label_ary, int w, int h, morph_label* label)
{
	int size = w * h;
	BYTE color = 255;

	BYTE* xkpre = malloc(size);
	BYTE* xk = malloc(size);
	BYTE* tempA = malloc(size);

	//init point set to 1
	memset(xkpre, 0, size);
	int index = label->k0y * w + label->k0x;
	xkpre[index] = color;

	//main loop
	int k = 1;
	do {
		//chg = X(k - 1) + B) and ori^c
		//tempA = X(k - 1) + B
		dilation_eightcon(xkpre, tempA);
		//xk = tempA intersection A = X(k - 1) + B and ori^c
		morph_sets_intersection(tempA, ori, xk, size);

		//check if equal
		if(memcmp(xkpre, xk, size)) {
			k++;
			memcpy(xkpre, xk, size);
		} else {
			break;
		}
	}while(k < 1000);
	printf("foreground label %d: iteration step is %d\n", label->label, k);

	morph_stat_for_con(xk, label_ary, w, h, label, color);

	free(xkpre);
	free(xk);
	free(tempA);

	return k;
}

/*
 * label all connected components of foreground
 *
 * input:
 *  h, w: image height and weight
 *  ori[h][w]: image data
 *
 *output:
 *  label_ary[h][w]: label map
 *  morph_label label[MORPH_END_LABEL]: fill it
 *
 *return value:
 *	number of connected components
 * */
int morph_label_foreground(BYTE* ori, BYTE* label_ary, int w, int h, morph_label* label)
{
	BYTE color;
	int i, j, size;

	if(!ori || !label)
		return 0;

	size = w * h;
	//find foreground color
	for(i = 0; i < size; i++) {
		if(ori[i]) {
			break;
		}
	}
	color = ori[i];
	BYTE* mirror = malloc(size);
	memset(mirror, 0, size);

	BYTE* chg = malloc(size);
	memcpy(chg, ori, size);
	memset(label_ary, 0, size);
	printf("\nstart to find all connected components area of this image\n\n");
	morph_label* cur;
	int knum = 0;
	int total_steps = 0, ksteps, run_flag = 1;
	BYTE cur_label = MORPH_START_LABEL;
	while(run_flag) {
		cur = &label[knum];
		//find first foreground color
		for(i = 0; i < size; i++) {
			if(chg[i]) {
				cur->k0y = i / w;
				cur->k0x = i % w;
				break;
			}
		}

		cur->label = cur_label;
		ksteps = morph_label_fore_con(chg, label_ary, w, h, cur);
		total_steps += ksteps;

		//copy this connected components from chg to mirror
		//and set related area of chg to 0(background color)
		for(j = 0; j < size; j++) {
			if(label_ary[j] == cur_label) {
				mirror[j] = chg[j];
				chg[j] = 0;
			}
		}

		//end loop when all chg[] is zero
		run_flag = 0;
		for(j = 0; j < size; j++) {
			if(chg[j]) {
				run_flag = 1;
				break;
			}
		}

		//handle next connected area
		cur_label++;
		knum++;
	}
	printf("total found %d connected area\n", knum);
	printf("total iteration steps is %d\n", total_steps);
	//mark end flag
	cur = &label[knum];
	cur->label = MORPH_END_LABEL;

	free(mirror);
	free(chg);

	return knum;
}

/* Morphological Reconstruction
 * dilation = 1: dilation
 * otherwise: erosion
 */
void morph_reconstruction(BYTE* markerF, BYTE* maskG, BYTE* resultR,
	int w, int h, int dilation)
{
	BYTE* F = markerF;
	BYTE* G = maskG;

	int size = w * h;
	BYTE* pre = malloc(size); //DG k
	BYTE* cur = resultR; //DG K + 1

	int k = 0;
	memcpy(pre, F, size); //DG0 = F
	//main loop
	while(k < 1000) {
		if(dilation) {
			dilation_eightcon(pre, cur);
			morph_sets_intersection(cur, G, cur, size);
		} else {
			erosion_eightcon(pre, cur);
			morph_sets_union(cur, G, cur, size);
		}
		if(!memcmp(pre, cur, size))
			break;
		memcpy(pre, cur, size);

		k++;
		if(k % 10 == 0)
			printf("reconstruction, iteration step k = %d\n", k);
	}
	printf("Morphological reconstruction iteration step is %d\n\n", k);

	free(pre);
}

/*
 * set border value according to boder_invert
 * fill hole: boder_invert = 1, border = 255 - border
 * border clear: boder_invert = 0, border value not change
 * not border: value = 0
 * */
void morph_get_marker(BYTE* ori, BYTE* chg, int w, int h, int boder_invert)
{
	int x, y, size, index;

	size = w * h;
	memset(chg, 0, size);
	for(y = 0; y < h; y++) {
		for(x = 0; x < w; x++) {
			if(y==0 || y==(h-1) || x==0 || x==(w-1)) {
				index = y * w + x;
				if(boder_invert)
					chg[index] = 255 - ori[index];
				else
					chg[index] = ori[index];
			}
		}//end for x
	}//end for y
}

///////////////////////////////////////////////////////////////////////////////
//gray-scale morphology

/* flat morphology structuring element
 * strel('disk', r)
 *
 * r > 0
 * dis: real r
 *
 * output: SE
 * call free(se) to free it
 */
BYTE* morph_strel_disk(int r, int* dis)
{
	BYTE* se = NULL;
	if(r <= 0)
		return NULL;

	int realR = r;
	int size, dim;
	switch(r) {
		case 1: //3 * 3
			realR = 1;
			break;
		case 2: //5 * 5
			realR = 2;
			break;
		case 3: //5 * 5
			realR = 2;
			break;
		default:
			realR = r - 1;
			break;
	}

	dim = 2 * realR + 1;
	size = dim * dim;
	se = malloc(size);
	if(r == 3) {
		memset(se, 1, size);
	} else {
		int x, y, a, rr, index;
		rr = r * r;
		//printf("r = %d\n", r);
		for(y = 0; y < dim; y++) {
			//printf("\n");
			for(x = 0; x < dim; x++) {
				a = (y - realR) * (y - realR) + (x - realR) * (x - realR);
				index = y * dim + x;
				if(a <= rr)
					se[index] = 1;
				else
					se[index] = 0;
				//printf("%d ", se[index]);
			} //end for x
		}//end for y
	}
	*dis = realR;
	return se;
}

/* A - B
 * w, h: image width and height
 * r: radius of SE
 * ori[h][w] to chg[h][w]
 * */
//f - B = min(f(x + s, y + t)) = min filter
void erosion_mem_gray_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	int x, y, p, q, k, xdis, ydis;
	int index, ii;
	BYTE min;
	BYTE* se = morph_strel(type, r, &xdis, &ydis);
	if(!se)
		return;

	for(y = 0; y < h; y++){
		for(x = 0; x < w; x++) {
			index = y * w + x;
			min = ori[index];
			k = 0;
			for(q = y - ydis; q <= y + ydis; q++) {
				for(p = x - xdis; p <= x + xdis; p++) {
					if(q < 0 || q >= h || p < 0 || p >= w) {
						k++;
						continue;
					}
					ii = q * w + p;
					if(ori[ii] < min && se[k])
						min = ori[ii];
					k++;
				}//end for p
			}//end for q
			chg[index] = min;
		}//end for x
	}//end for y

	free(se);
}

//A + B
//f + B = max(f(x - s, y - t)) = max filter
void dilation_mem_gray_common(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	int x, y, p, q, k, xdis, ydis;
	int index, ii;
	BYTE max;
	BYTE* se = morph_strel(type, r, &xdis, &ydis);
	if(!se)
		return;

	for(y = 0; y < h; y++){
		for(x = 0; x < w; x++) {
			index = y * w + x;
			max = ori[index];
			k = 0;
			for(q = y - ydis; q <= y + ydis; q++) {
				for(p = x - xdis; p <= x + xdis; p++) {
					if(q < 0 || q >= h || p < 0 || p >= w) {
						k++;
						continue;
					}
					ii = q * w + p;
					if(ori[ii] > max && se[k])
						max = ori[ii];
					k++;
				}//end for p
			}//end for q
			chg[index] = max;
		}//end for x
	}//end for y

	free(se);
}

void morph_open_common_gray(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	BYTE* temp = malloc(w * h);
	erosion_mem_gray_common(ori, temp, w, h, type, r);
	dilation_mem_gray_common(temp, chg, w, h, type, r);
	free(temp);
}

void morph_close_common_gray(BYTE* ori, BYTE* chg, int w, int h, SE_TYPE type, int r)
{
	BYTE* temp = malloc(w * h);
	dilation_mem_gray_common(ori, temp, w, h, type, r);
	erosion_mem_gray_common(temp, chg, w, h, type, r);
	free(temp);
}

//open first, and then close
void morph_smooth_gray(BYTE* ori, BYTE* chg, int w, int h, int r)
{
	BYTE* temp = malloc(w * h);
	morph_open_gray(ori, temp,  r);
	morph_close_gray(temp, chg, r);
	free(temp);
}

//(f + b) - (f - b)
void morph_gradient(BYTE* ori, BYTE* chg, int w, int h, int r)
{
	int size = w * h;
	BYTE* add = malloc(size);
	BYTE* sub = malloc(size);

	dilation_mem_gray_common(ori, add, w, h, SE_TYPE_SQUARE, r);
	erosion_mem_gray_common(ori, sub, w, h, SE_TYPE_SQUARE, r);
	mem_add(add, sub, chg, -1, w, h);

	free(add);
	free(sub);
}

//top-hat transformation
//f - open(f, b)
void morph_tophat(BYTE* ori, BYTE* chg, int w, int h, int r)
{
	morph_open_gray(ori, chg, r);
	mem_add(ori, chg, chg, -1, w, h);
}

//bottom-hat transformation
//close(f, b) - f
void morph_bothat(BYTE* ori, BYTE* chg, int w, int h, int r)
{
	morph_close_gray(ori, chg, r);
	mem_add(chg, ori, chg, -1, w, h);
}

//min = min(a,b)
void image_get_min(BYTE* a, BYTE* b, BYTE* min, int size)
{
	int i;
	for(i = 0; i < size; i++) {
		if(a[i] < b[i])
			min[i] = a[i];
		else
			min[i] = b[i];
	}
}

//max = max(a,b)
void image_get_max(BYTE* a, BYTE* b, BYTE* max, int size)
{
	int i;
	for(i = 0; i < size; i++) {
		if(a[i] < b[i])
			max[i] = b[i];
		else
			max[i] = a[i];
	}
}

//Morphological Reconstruction
// dilation = 1: dilation
// otherwise: erosion
void morph_recon_gray(BYTE* markerF, BYTE* maskG, BYTE* resultR,
	int w, int h, int dilation)
{
	BYTE* F = markerF;
	BYTE* G = maskG;

	int size = w * h;
	BYTE* pre = malloc(size); //DG k
	BYTE* cur = resultR; //DG K + 1

	int k = 0;
	memcpy(pre, F, size); //DG0 = F
	//main loop
	while(k < 1000) {
		if(dilation) {
			dilation_gray_eightcon(pre, cur);
			image_get_min(cur, G, cur, size);
		} else {
			erosion_gray_eightcon(pre, cur);
			image_get_max(cur, G, cur, size);
		}
		if(!memcmp(pre, cur, size))
			break;
		memcpy(pre, cur, size);

		k++;
		if(k % 10 == 0)
			printf("reconstruction, iteration step k = %d\n", k);
	}
	printf("Morphological reconstruction iteration step is %d\n\n", k);

	free(pre);
}


